import os 
if os.name == 'nt':
    os.chdir("C:\\federated_malware\\Project")
    print ("cwd",os.getcwd())#获得当前目录
    print ("工作目录",os.path.abspath('.'))#获得当前工作目录
    print ("工作目录",os.path.abspath(os.curdir))#获得当前工作目录
import logging
logger = logging.getLogger()
# 导入参数------------------------------------------------------
from fl_config import dataset_params,server_params,client_params

logger.info(dataset_params)
logger.info(server_params)
logger.info(client_params)


# 导入数据------------------------------------------------------
# 初始化各种类型的数据集。这里返回的是不同数据的索引，不能修改，否则会影响原来的数据。可以复用这些索引。因为创建dataset的时候一般是复制一份，而且一般也不会对原来的数据集进行修改。
from datasets.minist.data_load import dataManager as mdm
# from datasets.feminist.data_load import dataManager as fdm 
from datasets.maldroid.data_load import dataManager as maldm 

# 只需要能跑过就行，统计的返回值Counter符合要求即可.
# 确实，规定数据量不太合理。导致数据缺失，显然结果会不好。
# 所以现在完全采取互斥划分的五种方式。正好可以减少数据划分的情况。
# mdm_avg = mdm.allocate_data_avg(client_number=dataset_params["client_number"])
# mdm_ln = mdm.allocate_data_noniid_by_label_number(client_number=dataset_params["client_number"],class_total=10,class_counts=dataset_params["class_counts"])
# mdm_ld = mdm.allocate_data_noniid_by_label_distribution(client_number=dataset_params["client_number"],class_total=10,beta=dataset_params["beta"])
# mdm_fn = mdm.allocate_data_noniid_by_feature_noise(client_number=dataset_params["client_number"],bias_level=dataset_params["bias_level"])
# mdm_quantity = mdm.allocate_data_noniid_by_quantity(client_number=dataset_params["client_number"],beta=dataset_params["beta"])
# mdm_test = mdm.allocate_data_test()

# fdm_noniid = fdm.allocate_data_noniid()
# fdm_test = fdm.allocate_data_test()


maldm_avg = maldm.allocate_data_avg(client_number=dataset_params["client_number"])
maldm_ln = maldm.allocate_data_noniid_by_label_number(client_number=dataset_params["client_number"],class_total=5,class_counts=2)
maldm_ld = maldm.allocate_data_noniid_by_label_distribution(client_number=20,class_total=5,beta=0.4)
maldm_fn = maldm.allocate_data_noniid_by_feature_noise(client_number=20,bias_level=0.1)
maldm_quantity = maldm.allocate_data_noniid_by_quantity(client_number=20,beta=0.4)
maldm_test = maldm.allocate_data_test()


# 导入模型------------------------------------------------------
# from models.minist.dnn_model import Net
from models.maldroid.dnn_model import Net 
model = Net()
# param_list = model.parameters()

# 导入算法------------------------------------------------------
from trainers.fedbase.FedServerBase import FedServerBase
from trainers.fedbase.FedClientBase import FedClientBase
from trainers.fedavg.FedAvgServer import FedAvgServer
from trainers.fedavg.FedAvgClient import FedAvgClient
from trainers.fedida.FedIdaServer import FedIdaServer
from trainers.fedida.FedIdaClient import FedIdaClient
from trainers.fedprox.FedProxServer import FedProxServer
from trainers.fedprox.FedProxClient import FedProxClient
from trainers.fedscaffold.FedScaffoldServer import FedScaffoldServer
from trainers.fedscaffold.FedScaffoldClient import FedScaffoldClient

from trainers.fedamp.FedAmpServer import FedAmpServer
from trainers.fedamp.FedAmpClient import FedAmpClient
from trainers.fedper.FedPerServer import FedPerServer
from trainers.fedper.FedPerClient import FedPerClient
from trainers.fedmaml.FedMamlServer import FedMamlServer
from trainers.fedmaml.FedMamlClient import FedMamlClient
from trainers.fedreptile.FedReptileServer import FedReptileServer
from trainers.fedreptile.FedReptileClient import FedReptileClient
from trainers.fedmeta.FedMetaServer import FedMetaServer
from trainers.fedmeta.FedMetaClient import FedMetaClient

# 3个数据集 + 5中划分方法+9个算法
# 每个算法独立实验。每个数据集独立实验。3个数据集+5中划分方法，通过循环完成。
# ------------------------------------------------fedbase_information----------------------------------------
# train_worker_list = [FedClientBase(train_dataset,model) for train_dataset in mdm_avg]
# test_worker_list = [FedClientBase(test_dataset,model) for test_dataset in mdm_test]
# server = FedServerBase((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()
# server.save_state('./result/mdm_avg_fedbase')
# ------------------------------------------------fedbase----------------------------------------
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedClientBase(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedClientBase(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedServerBase((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# 创建多个客户端和服务器。然后运行base联邦学习。noniid数据
# train_worker_list = [FedClientBase(train_dataset,model) for train_dataset in mdm_avg]
# test_worker_list = [FedClientBase(test_dataset,model) for test_dataset in mdm_test]
# server = FedServerBase((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# ------------------------------------------------fedavg----------------------------------------
# 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# train_dataset_list = [fdm_noniid]
# test_dataset_list = [fdm_test]
# dataset_name = ['fdm_noniid']

train_dataset_list = [maldm_avg,maldm_ln,maldm_ld,maldm_fn,maldm_quantity]
test_dataset_list = [maldm_test,maldm_test,maldm_test,maldm_test,maldm_test]
dataset_name = ['maldm_avg','maldm_ln','maldm_ld','maldm_fn','maldm_quantity']

# base_name = './results/'
# algorithm_name = '_fedavg'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedAvgClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedAvgClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedAvgServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)


# ------------------------------------------------fedida----------------------------------------
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedIdaClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedIdaClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedIdaServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()


# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']


# base_name = './results/'
# algorithm_name = '_fedida'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedIdaClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedIdaClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedIdaServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)

# ------------------------------------------------fedprox----------------------------------------
# client_params["prox_mu"]=0.1
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedProxClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedProxClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedProxServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# client_params["prox_mu"]=0.1

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# base_name = './results/'
# algorithm_name = '_fedprox'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedProxClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedProxClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedProxServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)


# ------------------------------------------------fedscaffold----------------------------------------
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedScaffoldClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedScaffoldClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedScaffoldServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()


# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']

# base_name = './results/'
# algorithm_name = '_fedscaffold'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedScaffoldClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedScaffoldClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedScaffoldServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)


# ------------------------------------------------fedamp----------------------------------------
# # 添加fedamp相关的控制参数
# server_params["ak"]=0
# server_params["sigma"]=10
# client_params["amp_lambda"]=0.0
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedAmpClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedAmpClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedAmpServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# train_dataset_list = [mdm_avg,mdm_avg,mdm_avg,mdm_avg,mdm_avg]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_avg','mdm_avg','mdm_avg','mdm_avg']
# ak = [0.8,0.8,0.9,0.9,0.9]
# sigma = [1,10,1,10,0.2]

# 添加fedamp相关的控制参数
# 发现0.8-1,0.8-10,0.9-10这三个参数非常好
# server_params["ak"]=0.9
# server_params["sigma"]=10
# client_params["amp_lambda"]=0.1
# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']

# 在mdm_avg下测一组变量

# base_name = './results/'
# algorithm_name = '_fedamp'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedAmpClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedAmpClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedAmpServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)

# ------------------------------------------------fedper----------------------------------------
# server_params["base_layer"]=4
# client_params["base_layer"]=4
# client_params["support_ratio"]=0.7
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedPerClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedPerClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedPerServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# server_params["base_layer"]=4
# client_params["base_layer"]=4
# client_params["support_ratio"]=0.7
# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']

# base_name = './results/'
# algorithm_name = '_fedper'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedPerClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedPerClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedPerServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix) 
    

# ------------------------------------------------fedmaml----------------------------------------
# client_params["support_ratio"]=0.7
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedMamlClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedMamlClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedMamlServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()


# client_params["support_ratio"]=0.7

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# base_name = './results/'
# algorithm_name = '_fedmaml'
# end_tree = '_.pt'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedMamlClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedMamlClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedMamlServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)


# ------------------------------------------------fedreptile----------------------------------------
# client_params["support_ratio"]=0.7
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedReptileClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedReptileClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedReptileServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# 创建多个客户端和服务器。然后运行base联邦学习。noniid数据
# client_params["support_ratio"]=0.7
# server_params['outer_lr']=0.8
# client_params['outer_lr']=0.8

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# base_name = './results/'
# algorithm_name = '_fedreptile'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedReptileClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedReptileClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedReptileServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)


# ------------------------------------------------fedmeta----------------------------------------
# server_params["base_layer"]=4
# client_params["base_layer"]=4
# client_params["support_ratio"]=0.7
# # 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
# train_worker_list = [FedMetaClient(train_dataset,model) for train_dataset in train_dataset_list_iid]
# test_worker_list = [FedMetaClient(test_dataset,model) for test_dataset in test_dataset_list]
# server = FedMetaServer((train_worker_list,test_worker_list),model,server_params,client_params)
# server.start_federated()

# 创建多个客户端和服务器。然后运行base联邦学习。noniid数据
server_params["base_layer"]=4
client_params["base_layer"]=4
client_params["support_ratio"]=0.7

# reptile元学习中会用到参数
server_params['outer_lr']=0.8
client_params['outer_lr']=0.8

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
base_name = './results/'
algorithm_name = '_fedmeta'
suffix = '_.pt'

for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
    train_worker_list = [FedMetaClient(train_dataset,model) for train_dataset in train_data]
    test_worker_list = [FedMetaClient(test_dataset,model) for test_dataset in test_dataset]
    server = FedMetaServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
    logger.info("本次测试的数据集的名字:{}".format(name))
    logger.info("本次测试的算法的名字:{}".format(algorithm_name))
    # server.start_federated_with_personality_layer()
    server.start_federated_with_reptile_adjust()
    server.save_state(base_name+name+algorithm_name+suffix)


for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
    train_worker_list = [FedMetaClient(train_dataset,model) for train_dataset in train_data]
    test_worker_list = [FedMetaClient(test_dataset,model) for test_dataset in test_dataset]
    server = FedMetaServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
    logger.info("本次测试的数据集的名字:{}".format(name))
    logger.info("本次测试的算法的名字:{}".format(algorithm_name))
    server.start_federated_with_personality_layer()
    # server.start_federated_with_reptile_adjust()
    server.save_state(base_name+name+algorithm_name+"layer"+suffix)